import os
import time
from concurrent.futures import ThreadPoolExecutor

import requests
from Crypto.Cipher import AES


def get_m3u8(url):
    res = requests.get(url)
    with open("../数据暂存/m3u8.txt", "wb") as f:
        f.write(res.content)
    print("m3u8文件下载完成")


def get_key(url):
    with open("../数据暂存/m3u8.txt", "r", encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if line.startswith('#EXT-X-KEY'):
                key_url = line.split('"')[1]
                print("key_url:" + key_url)
                break
    res = requests.get(url + "/" + key_url)
    res.encoding = 'utf-8'
    return res.content


def dec_ts(key, data):
    """AES解密视频"""
    aes = AES.new(key, mode=AES.MODE_CBC)
    data = aes.decrypt(data)
    return data


def transform(input_str):
    # 提取输入字符串中的所有数字
    digits = ''.join(filter(str.isdigit, input_str))
    # 将提取出的数字补充成三位数
    padded_number = digits.zfill(3)
    return padded_number


def down_ts(url, key):
    """下载单个ts文件"""
    ts_name0 = url[-6:-3]
    ts_name = transform(ts_name0) + ".ts"
    while True:
        try:
            res = requests.get(url, headers=head, timeout=20)
            with open(f"../数据暂存/{ts_name}", "wb") as f:
                data = res.content
                data = dec_ts(key, data)
                f.write(data)
            print(f"{ts_name}下载完成")
            # 记录已下载的文件名称
            global lst
            lst.append(ts_name)
            break
        except Exception as e:
            print(f"{ts_name}下载出现异常", e)
            time.sleep(5)


def start_down_ts(key):
    list1 = []
    with open("../数据暂存/m3u8.txt", "r", encoding="utf-8") as f:
        line: str
        for line in f:
            if line.startswith("#") or len(line) < 2:
                continue
            ts_url = line.strip()
            list1.append(ts_url)
    print(len(list1))
    with ThreadPoolExecutor(25) as t:
        for i in list1:
            t.submit(down_ts, i, key)
    print("全部下载完成")


def merge(name):
    """合并ts文件，生成完整的mp4文件"""
    download_path = r'E:\T\Videos\SpiderDownload'
    download_name = os.path.join(download_path, name)
    print(download_name)
    # 获取所有ts文件列表
    global lst
    lst.sort()
    os.chdir("../数据暂存")
    os.system(f"copy /b *.ts {download_name}")
    print(f"合并成功，文件已存放在{download_path}")
    # 删除暂存数据
    # os.system("del *.ts")
    os.system("del *.txt")
    print("======Done=====")


def main(url, name):
    # 获取m3u8
    get_m3u8(url)
    # 获取密钥
    main_url = url.rsplit("/", 1)[0]
    key = get_key(main_url)
    # 下载源文件
    start_down_ts(key)
    # 视频整合
    merge(name)


if __name__ == '__main__':
    # m3u8文件网址
    url1 = "https://hnzy.bfvvs.com/play/wdLYwx4e/index.m3u8"
    url2 = "https://hnzy.bfvvs.com/play/qaQNBDYb/index.m3u8"
    url3 = "https://hnzy.bfvvs.com/play/QbYQKGYb/index.m3u8"
    url4 = "https://hnzy.bfvvs.com/play/qaQAW8Ya/index.m3u8"
    url5 = "https://v.gsuus.com/play/9av2q18a/index.m3u8"
    head = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36 Edg/126.0.0.0"}
    url = input("请输入下载视频的m3u8地址：\n")
    name = input("请输入下载的视频名称(请不要输入空格)：\n") + ".mp4"
    # 定义一个列表，记录ts文件名称
    lst = []
    # 该列表记载两次异常的请求
    t1 = time.time()
    main(url, name)
    t2 = time.time()
    print(f"总耗时{t2 - t1}s")
